#!/usr/bin/env python3
"""
数据拆分插件Pipeline演示
展示如何在实际数据处理流程中使用数据拆分插件
"""

import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.plugins.data_splitter.data_splitter import DataSplitter


def create_sample_dataset():
    """创建一个模拟的数据集"""
    return [
        {"customer_id": 1, "age": 25, "income": 3000, "credit_score": 650, "loan_amount": 10000, "approved": 1},
        {"customer_id": 2, "age": 35, "income": 5000, "credit_score": 720, "loan_amount": 20000, "approved": 1},
        {"customer_id": 3, "age": 45, "income": 7000, "credit_score": 680, "loan_amount": 15000, "approved": 0},
        {"customer_id": 4, "age": 28, "income": 3500, "credit_score": 590, "loan_amount": 8000, "approved": 0},
        {"customer_id": 5, "age": 42, "income": 6000, "credit_score": 750, "loan_amount": 25000, "approved": 1},
        {"customer_id": 6, "age": 31, "income": 4500, "credit_score": 620, "loan_amount": 12000, "approved": 1},
        {"customer_id": 7, "age": 38, "income": 5500, "credit_score": 690, "loan_amount": 18000, "approved": 1},
        {"customer_id": 8, "age": 29, "income": 3200, "credit_score": 580, "loan_amount": 9000, "approved": 0},
        {"customer_id": 9, "age": 33, "income": 4800, "credit_score": 710, "loan_amount": 16000, "approved": 1},
        {"customer_id": 10, "age": 27, "income": 3800, "credit_score": 640, "loan_amount": 11000, "approved": 0},
        {"customer_id": 11, "age": 36, "income": 5200, "credit_score": 730, "loan_amount": 22000, "approved": 1},
        {"customer_id": 12, "age": 41, "income": 6500, "credit_score": 700, "loan_amount": 19000, "approved": 1},
        {"customer_id": 13, "age": 26, "income": 3100, "credit_score": 600, "loan_amount": 8500, "approved": 0},
        {"customer_id": 14, "age": 34, "income": 4200, "credit_score": 660, "loan_amount": 13000, "approved": 1},
        {"customer_id": 15, "age": 39, "income": 5800, "credit_score": 740, "loan_amount": 21000, "approved": 1},
        {"customer_id": 16, "age": 30, "income": 4000, "credit_score": 630, "loan_amount": 10500, "approved": 0},
        {"customer_id": 17, "age": 37, "income": 5100, "credit_score": 680, "loan_amount": 17000, "approved": 1},
        {"customer_id": 18, "age": 32, "income": 4600, "credit_score": 650, "loan_amount": 14000, "approved": 1},
        {"customer_id": 19, "age": 43, "income": 6200, "credit_score": 760, "loan_amount": 24000, "approved": 1},
        {"customer_id": 20, "age": 24, "income": 2800, "credit_score": 570, "loan_amount": 7500, "approved": 0}
    ]


def demonstrate_train_test_split():
    """演示训练集/测试集拆分"""
    print("=== 机器学习场景：训练集/测试集拆分 ===")
    
    # 创建数据拆分插件
    config = {
        "config": {
            "split_strategy": "ratio",
            "ratios": [0.8, 0.2],  # 80%训练集，20%测试集
            "random_seed": 42,     # 固定随机种子
            "shuffle": True        # 打乱数据
        }
    }
    splitter = DataSplitter("train_test_splitter", config)
    
    # 模拟从数据加载器获取的数据
    dataset = create_sample_dataset()
    pipeline_data = {"data_loader": {"raw_data": dataset}}
    
    # 执行拆分
    result = splitter.execute(pipeline_data)
    
    if result["split_data"]:
        train_data, test_data = result["split_data"]
        
        print(f"原始数据集: {result['original_count']} 条记录")
        print(f"训练集: {len(train_data)} 条记录 ({len(train_data)/result['original_count']*100:.1f}%)")
        print(f"测试集: {len(test_data)} 条记录 ({len(test_data)/result['original_count']*100:.1f}%)")
        
        # 显示训练集和测试集的统计信息
        train_approved = sum(1 for record in train_data if record["approved"] == 1)
        test_approved = sum(1 for record in test_data if record["approved"] == 1)
        
        print(f"训练集批准率: {train_approved/len(train_data)*100:.1f}%")
        print(f"测试集批准率: {test_approved/len(test_data)*100:.1f}%")
        
        print("训练集前3条记录:")
        for i, record in enumerate(train_data[:3]):
            print(f"  {i+1}. 客户{record['customer_id']}: 年龄{record['age']}, 收入{record['income']}, 批准{'是' if record['approved'] else '否'}")
        
        print("测试集前3条记录:")
        for i, record in enumerate(test_data[:3]):
            print(f"  {i+1}. 客户{record['customer_id']}: 年龄{record['age']}, 收入{record['income']}, 批准{'是' if record['approved'] else '否'}")
    
    print()


def demonstrate_age_group_split():
    """演示按年龄组拆分"""
    print("=== 业务分析场景：按年龄组拆分客户 ===")
    
    # 创建数据拆分插件 - 按年龄条件拆分
    config = {
        "config": {
            "split_strategy": "condition",
            "condition_field": "age",
            "condition_value": 35,
            "condition_operator": ">="  # 年龄>=35的为一组，其他的为另一组
        }
    }
    splitter = DataSplitter("age_group_splitter", config)
    
    # 执行拆分
    dataset = create_sample_dataset()
    result = splitter.execute(dataset)  # 直接使用列表格式
    
    if result["split_data"]:
        senior_group, junior_group = result["split_data"]
        
        print(f"年龄>=35的客户: {len(senior_group)} 人")
        print(f"年龄<35的客户: {len(junior_group)} 人")
        
        # 计算每组的平均收入、信用评分等统计信息
        senior_avg_income = sum(record["income"] for record in senior_group) / len(senior_group)
        junior_avg_income = sum(record["income"] for record in junior_group) / len(junior_group)
        
        senior_avg_credit = sum(record["credit_score"] for record in senior_group) / len(senior_group)
        junior_avg_credit = sum(record["credit_score"] for record in junior_group) / len(junior_group)
        
        senior_approval_rate = sum(1 for record in senior_group if record["approved"] == 1) / len(senior_group) * 100
        junior_approval_rate = sum(1 for record in junior_group if record["approved"] == 1) / len(junior_group) * 100
        
        print(f"\n年龄>=35组统计:")
        print(f"  平均收入: {senior_avg_income:.0f}")
        print(f"  平均信用评分: {senior_avg_credit:.0f}")
        print(f"  批准率: {senior_approval_rate:.1f}%")
        
        print(f"\n年龄<35组统计:")
        print(f"  平均收入: {junior_avg_income:.0f}")
        print(f"  平均信用评分: {junior_avg_credit:.0f}")
        print(f"  批准率: {junior_approval_rate:.1f}%")
    
    print()


def demonstrate_credit_score_split():
    """演示按信用评分拆分"""
    print("=== 风险管理场景：按信用评分拆分 ===")
    
    # 创建多个数据拆分插件，实现多级拆分
    
    # 第一步：按信用评分粗略分组
    high_credit_config = {
        "config": {
            "split_strategy": "condition",
            "condition_field": "credit_score",
            "condition_value": 700,
            "condition_operator": ">="  # 高信用评分
        }
    }
    high_credit_splitter = DataSplitter("high_credit_splitter", high_credit_config)
    
    # 执行第一级拆分
    dataset = create_sample_dataset()
    result1 = high_credit_splitter.execute(dataset)
    
    if result1["split_data"]:
        high_credit_group, low_credit_group = result1["split_data"]
        
        print(f"高信用评分组(>=700): {len(high_credit_group)} 人")
        print(f"低信用评分组(<700): {len(low_credit_group)} 人")
        
        # 第二步：对高信用评分组进一步按收入拆分
        high_income_config = {
            "config": {
                "split_strategy": "condition",
                "condition_field": "income",
                "condition_value": 5000,
                "condition_operator": ">="  # 高收入
            }
        }
        high_income_splitter = DataSplitter("high_income_splitter", high_income_config)
        
        result2 = high_income_splitter.execute(high_credit_group)
        
        if result2["split_data"]:
            premium_customers, standard_customers = result2["split_data"]
            print(f"\n高信用评分组进一步拆分:")
            print(f"  高信用+高收入(优质客户): {len(premium_customers)} 人")
            print(f"  高信用+低收入(标准客户): {len(standard_customers)} 人")
        
        # 对低信用评分组也进行类似拆分
        low_income_config = {
            "config": {
                "split_strategy": "condition",
                "condition_field": "income",
                "condition_value": 4000,
                "condition_operator": ">="  # 收入>=4000
            }
        }
        low_income_splitter = DataSplitter("low_income_splitter", low_income_config)
        
        result3 = low_income_splitter.execute(low_credit_group)
        
        if result3["split_data"]:
            recoverable_customers, high_risk_customers = result3["split_data"]
            print(f"\n低信用评分组进一步拆分:")
            print(f"  低信用+较高收入(可挽回客户): {len(recoverable_customers)} 人")
            print(f"  低信用+低收入(高风险客户): {len(high_risk_customers)} 人")
    
    print()


def demonstrate_batch_processing():
    """演示批量处理场景"""
    print("=== 批量处理场景：按批次拆分处理 ===")
    
    # 创建数据拆分插件 - 按固定数量拆分
    config = {
        "config": {
            "split_strategy": "count",
            "counts": [5, 5, 5, 5],  # 每批5条记录，共4批
            "remainder_strategy": "distribute"  # 剩余数据分配到前几批
        }
    }
    batch_splitter = DataSplitter("batch_processor", config)
    
    # 执行拆分
    dataset = create_sample_dataset()
    result = batch_splitter.execute(dataset)
    
    if result["split_data"]:
        batches = result["split_data"]
        
        print(f"数据集被拆分为 {len(batches)} 个批次:")
        for i, batch in enumerate(batches):
            print(f"\n批次 {i+1}: {len(batch)} 条记录")
            
            # 模拟批处理统计
            total_income = sum(record["income"] for record in batch)
            avg_credit_score = sum(record["credit_score"] for record in batch) / len(batch)
            approved_count = sum(1 for record in batch if record["approved"] == 1)
            
            print(f"  总收入: {total_income}")
            print(f"  平均信用评分: {avg_credit_score:.0f}")
            print(f"  批准数量: {approved_count}/{len(batch)}")
            print(f"  客户ID: {[record['customer_id'] for record in batch]}")
    
    print()


def demonstrate_k_fold_validation():
    """演示K折交叉验证场景"""
    print("=== 机器学习场景：K折交叉验证 ===")
    
    # 创建数据拆分插件 - 实现K折交叉验证
    k = 5  # 5折交叉验证
    fold_size = len(create_sample_dataset()) // k
    
    config = {
        "config": {
            "split_strategy": "count",
            "counts": [fold_size] * k,  # k个大小相等的折
            "remainder_strategy": "distribute"  # 分配剩余数据
        }
    }
    k_fold_splitter = DataSplitter("k_fold_splitter", config)
    
    # 执行拆分
    dataset = create_sample_dataset()
    result = k_fold_splitter.execute(dataset)
    
    if result["split_data"]:
        folds = result["split_data"]
        
        print(f"{k}折交叉验证数据集拆分:")
        
        # 模拟交叉验证过程
        for fold_idx in range(k):
            print(f"\n第 {fold_idx + 1} 轮验证:")
            
            # 当前fold作为验证集
            validation_set = folds[fold_idx]
            
            # 其他folds作为训练集
            training_set = []
            for i, fold in enumerate(folds):
                if i != fold_idx:
                    training_set.extend(fold)
            
            print(f"  训练集: {len(training_set)} 条记录")
            print(f"  验证集: {len(validation_set)} 条记录")
            
            # 模拟模型评估
            training_approval_rate = sum(1 for record in training_set if record["approved"] == 1) / len(training_set) * 100
            validation_approval_rate = sum(1 for record in validation_set if record["approved"] == 1) / len(validation_set) * 100
            
            print(f"  训练集批准率: {training_approval_rate:.1f}%")
            print(f"  验证集批准率: {validation_approval_rate:.1f}%")
    
    print()


def demonstrate_stratified_sampling():
    """演示分层抽样场景"""
    print("=== 统计学场景：分层抽样 ===")
    
    # 按批准状态分层，保持原始比例
    approval_splitter = DataSplitter("approval_splitter", {
        "config": {
            "split_strategy": "field_value",
            "split_field": "approved",
            "max_groups": 2
        }
    })
    
    dataset = create_sample_dataset()
    result = approval_splitter.execute(dataset)
    
    if result["split_data"]:
        approved_group, rejected_group = result["split_data"]
        
        print(f"按批准状态分层:")
        print(f"  已批准组: {len(approved_group)} 人 ({len(approved_group)/len(dataset)*100:.1f}%)")
        print(f"  被拒绝组: {len(rejected_group)} 人 ({len(rejected_group)/len(dataset)*100:.1f}%)")
        
        # 从每层中按比例抽样
        sample_ratios = [0.5, 0.5]  # 从每层中抽取50%
        
        approved_sampler = DataSplitter("approved_sampler", {
            "config": {
                "split_strategy": "ratio",
                "ratios": [sample_ratios[0], 1-sample_ratios[0]],
                "random_seed": 42
            }
        })
        
        rejected_sampler = DataSplitter("rejected_sampler", {
            "config": {
                "split_strategy": "ratio",
                "ratios": [sample_ratios[1], 1-sample_ratios[1]],
                "random_seed": 42
            }
        })
        
        approved_sample_result = approved_sampler.execute(approved_group)
        rejected_sample_result = rejected_sampler.execute(rejected_group)
        
        if approved_sample_result["split_data"] and rejected_sample_result["split_data"]:
            approved_sample = approved_sample_result["split_data"][0]
            rejected_sample = rejected_sample_result["split_data"][0]
            
            print(f"\n分层抽样结果:")
            print(f"  已批准样本: {len(approved_sample)} 人")
            print(f"  被拒绝样本: {len(rejected_sample)} 人")
            print(f"  总样本: {len(approved_sample) + len(rejected_sample)} 人")
            print(f"  样本比例: 批准{(len(approved_sample)/(len(approved_sample)+len(rejected_sample))*100):.1f}%, 拒绝{(len(rejected_sample)/(len(approved_sample)+len(rejected_sample))*100):.1f}%")
    
    print()


if __name__ == "__main__":
    print("数据拆分插件Pipeline演示")
    print("=" * 60)
    print()
    
    # 运行各种场景演示
    demonstrate_train_test_split()
    demonstrate_age_group_split()
    demonstrate_credit_score_split()
    demonstrate_batch_processing()
    demonstrate_k_fold_validation()
    demonstrate_stratified_sampling()
    
    print("所有演示完成！")
    print("\n数据拆分插件支持多种业务场景:")
    print("- 机器学习: 训练集/测试集拆分、K折交叉验证")
    print("- 业务分析: 客户分群、风险分层")
    print("- 批量处理: 大数据分批处理")
    print("- 统计分析: 分层抽样、随机抽样")